# Title     : TODO
# Objective : TODO
# Created by: Administrator
# Created on: 2019/6/18

library(plyr)
library(tidyr)
library(gridExtra)
library(ggpubr)
library(egg)
library(lemon)
library(optparse)
library(ComplexHeatmap)
library(circlize)
library(tidyverse)

createWhenNoExist <- function(f) {
  !dir.exists(f) && dir.create(f)
}

options(digits = 3)

option_list <- list(
  make_option("--i", default = "AllMet.csv", type = "character", help = "metabolite data file"),
  make_option("--g", default = "SampleInfo.csv", type = "character", help = "sample group file"),
  make_option("--sc", default = "sample_color.txt", type = "character", help = "sample color file"),
  make_option("--mc", default = "meta_color.txt", type = "character", help = "metabolite color file")
)
opt <- parse_args(OptionParser(option_list = option_list))

sampleInfo <- read.csv(opt$g, header = T, stringsAsFactors = F) %>%
  select(c("SampleID", "ClassNote")) %>%
  mutate(ClassNote = as.character(ClassNote)) %>%
  arrange(ClassNote)

str(sampleInfo)

parent <- "./"
scoreFileName <- paste0(parent, "/Markers_Z_Score.csv")

if (!file.exists(scoreFileName)) {
  quit(status = 0)
}

scoreData <- read.csv(scoreFileName, header = T, check.names = F, row.names = 1)

scoreGroupData <- scoreData %>%
  rownames_to_column("SampleID") %>%
  left_join(sampleInfo, by = c("SampleID"))

metaColDf <- read_tsv(opt$mc) %>%
  select(c("Class", "col"))
metaCols <- metaColDf %>%
  deframe()

metaboliteInfo <- read.csv(opt$i, header = T, stringsAsFactors = F, check.names = F) %>%
  select(c("Class", "Metabolite")) %>%
  filter(Metabolite %in% colnames(scoreData))

metaboliteInfo

fileName <- paste0(parent, "/Markers_Z_Score_Heatmap_ordered_by_Class_with_Name.pdf")

tmpPlotData <- scoreData %>%
  t() %>%
  as.data.frame() %>%
  rownames_to_column("Metabolite") %>%
  left_join(metaboliteInfo, by = "Metabolite") %>%
  arrange(Class)

sortClass <- unique(tmpPlotData$Class) %>%
  map(function(eachClass) {
    plotData <- tmpPlotData %>%
      filter(Class == eachClass) %>%
      select(-c("Class")) %>%
      column_to_rownames("Metabolite")
    sortLabels <- if (nrow(plotData) > 1) {
      meta_dist <- dist(as.matrix(1 - (plotData)), method = "euclidean")
      hc_meta <- hclust(meta_dist, method = "complete")
      hc_meta$label[hc_meta$order]
    } else {
      rownames(plotData)
    }
    sortLabels
  }) %>%
  flatten_chr()

plotData <- tmpPlotData %>%
  mutate(Class = factor(Class, levels = sortClass)) %>%
  arrange(Class) %>%
  select(-c("Class")) %>%
  column_to_rownames("Metabolite")


if (nrow(plotData) < 2) {
  quit(status = 0)
}

width <- 7 + max(0, nrow(sampleInfo) - 20) * 0.1
rowNum <- nrow(plotData)
height <- max(2 + (rowNum - 10) * 0.1, 2) + 2

wHRate <- width / height
inchPx <- 72 / 2.54
maxAreaPx <- 10000
if (width * height * inchPx > maxAreaPx) {
  height <- sqrt(maxAreaPx / inchPx / wHRate)
  width <- wHRate * height
}

pdf(fileName, width = width, height = height)

sampleColDf <- read.csv(opt$sc, header = T, stringsAsFactors = F, comment.char = "") %>%
  select(c("ClassNote", "col"))
sampleCols <- sampleColDf %>%
  deframe()

colHa <- HeatmapAnnotation(ClassNote = scoreGroupData$ClassNote, col = list(ClassNote = sampleCols),
                           show_annotation_name = F, show_legend = F, simple_anno_size = unit(0.25, "cm")
)

colors <- rev(colorRampPalette(RColorBrewer::brewer.pal(10, "RdBu"))(256))

uniqSuperClass <- unique(metaboliteInfo$Class)
sortSuperClass <- uniqSuperClass %>%
  discard(~.x %in% c("Unknown")) %>%
  sort()
if (length(uniqSuperClass) != length(sortSuperClass)) {
  sortSuperClass <- sortSuperClass %>%
    c("Unknown")
}


sortMetaCols <- metaCols[sortSuperClass]

sortRowLgd <- Legend(at = sortSuperClass, legend_gp = gpar(fill = sortMetaCols), title = "Class",
                     labels_gp = gpar(fontsize = 12, fontfamily = "Times"), title_gp = gpar(fontsize = 12, fontfamily = "Times"))

legendMinV <- -3
legendMaxV<-3
colF <- colorRamp2(seq(legendMinV, legendMaxV, length.out = 256), colors)

zScoreLgd <- Legend(col_fun = colF, at = seq(legendMinV, legendMaxV, 1), title = "Z-Score",
                    title_gp = gpar(fontsize = 12, fontfamily = "Times"), grid_height = unit(5, "mm"),
                    labels_gp = gpar(fontsize = 12, fontfamily = "Times"))

pd <- packLegend(list = list(zScoreLgd, sortRowLgd), row_gap = unit(0.5, "cm"))

uniqClassNote <- unique(scoreGroupData$ClassNote)
inSampleCol <- sampleColDf %>%
  arrange(factor(ClassNote, levels = uniqClassNote))

colLgd <- Legend(at = inSampleCol$ClassNote, legend_gp = gpar(fill = inSampleCol$col), title = "",
                 labels_gp = gpar(fontsize = 12, fontfamily = "Times"), title_gp = gpar(fontsize = 12, fontfamily = "Times"), nrow = 1)

rowAnnoNames <- rownames(plotData)

needMetaboliteInfo <- metaboliteInfo %>%
  arrange(factor(Metabolite, levels = rowAnnoNames))

rowHa <- rowAnnotation(Class = needMetaboliteInfo$Class, col = list(Class = metaCols),
                       annotation_name_gp = gpar(fontsize = 6, fontfamily = "Times"), show_annotation_name = F, show_legend = F
)

htList <- Heatmap(plotData, col = colF, show_column_names = T, cluster_rows = F, cluster_columns = F,
                  name = "Z-Score", top_annotation = colHa, right_annotation = rowHa,
                  row_names_gp = gpar(fontsize = 6, fontfamily = "Times"), show_row_names = T,
                  column_names_gp = gpar(fontsize = 9, fontfamily = "Times"), show_heatmap_legend = F
)

draw(htList,
     padding = unit(c(0.5, 0.5, 0.5, 0.5), "cm"),
     annotation_legend_list = pd, heatmap_legend_list = list(colLgd), heatmap_legend_side = "top"
)

dev.off()



